import numpy as np
import matplotlib.pyplot as plt
import scipy

"""dct解泊松方程测试程序"""

def dct2(x):
    return scipy.fftpack.dct(scipy.fftpack.dct(x.T, norm='ortho').T, norm='ortho')

# 二维离散余弦反变换
def idct2(x):
    return scipy.fftpack.idct(scipy.fftpack.idct(x.T, norm='ortho').T, norm='ortho')

# 参数设置
Nx = 32
Ny = 32

# 网格间距
dx = 2 / Nx
dy = 2 / Ny

# 创建网格
x = np.arange(-1 + dx / 2, 1 - dx / 2 + dx/10, dx)
y = x = np.arange(-1 + dy / 2, 1 - dy / 2 + dy/10, dy)

X, Y = np.meshgrid(x, y, indexing='ij')  # 使用 'ij' 索引以匹配 MATLAB 的行为

# 修改波数
kx = np.arange(Nx)
ky = np.arange(Ny)
mwx = -(np.pi * kx / Nx)**2 / dx**2
mwy = -(np.pi * ky / Ny)**2 / dy**2

# 2D DCT of f (Right hand side)
f = -2 * np.pi**2 * np.cos(np.pi * X) * np.cos(np.pi * Y)
analy= np.cos(np.pi * X) * np.cos(np.pi * Y)

fhat = dct2(f)  # 使用 DCT-II 类型

# 构造 MWX 和 MWY 的网格
MWX, MWY = np.meshgrid(mwx, mwy, indexing='ij')

# 解方程
uhat = fhat / (MWX + MWY)

# 由于 uhat(0,0) 可能未定义，我们将其设置为 0
uhat[0, 0] = 0

# 逆 2D DCT
u = idct2(uhat)

# 绘制结果



# 计算误差
u_numeric=u
u_analytic=analy
error = np.abs(u_numeric - u_analytic)

# 绘制数值解、解析解和误差
plt.figure(figsize=(15, 5))

# 数值解
plt.subplot(1, 3, 1)
plt.title("Numerical Solution")
plt.imshow(u_numeric, extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 精确解
plt.subplot(1, 3, 2)
plt.title("Analytical Solution")
plt.imshow(u_analytic, extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 误差
plt.subplot(1, 3, 3)
plt.title("Error between Numerical and Analytical Solutions")
plt.imshow(error, extent=[-1, 1, -1, 1], origin='lower', cmap='coolwarm')
plt.colorbar(label='Error')

plt.tight_layout()
plt.show()